# 这是一个用于生成高质量问答对数据的脚本
# 实现了类似 https://hf-mirror.com/learn/cookbook/fine_tune_chatbot_docs_synthetic 中的方法

import os
import argparse
import asyncio
from data_generator import SyntheticDataGenerator


def create_sample_document():
    """
    创建一个示例文档用于测试
    """
    sample_doc = """
机器学习简介
===============================

机器学习是人工智能的一个领域，专注于开发能够让计算机在没有明确指令的情况下执行任务的算法和统计模型。
它们依靠模式和推理来做出决策。

机器学习的类型
-------------------------

1. 监督学习：在监督学习中，模型在标记的数据上进行训练。目标是学习从输入到输出的映射。
   常见算法包括线性回归、逻辑回归和支持向量机。

2. 无监督学习：在无监督学习中，模型被给予数据，但没有明确说明如何处理这些数据。
   目标是识别数据中的模式或结构。常见算法包括k均值聚类和主成分分析。

3. 强化学习：在强化学习中，智能体通过在环境中执行动作来学习做出决策，以最大化某种累积奖励的概念。
   这种方法通常用于机器人技术、游戏和导航。

机器学习中的关键概念
--------------------------------

- 特征：这些是被观察现象的各个可测量属性或特征。例如，在关于房屋的数据集中，
  特征可能包括大小、卧室数量和位置。

- 模型：对现实世界过程的数学表示。在机器学习中，模型被训练来对数据进行预测或识别模式。

- 训练：教授机器学习模型基于数据进行预测或决策的过程。

- 测试：在新的、未见过的数据上评估训练模型性能的过程。

机器学习的应用
--------------------------------

机器学习在各个行业都有众多应用：

1. 医疗保健：诊断、药物发现、医学影像
2. 金融：欺诈检测、算法交易、信用评分
3. 零售：推荐系统、库存管理、客户细分
4. 交通：自动驾驶汽车、路线优化、交通预测
5. 自然语言处理：语言翻译、情感分析、聊天机器人

机器学习面临的挑战
------------------------------

尽管机器学习具有潜力，但它仍面临几个挑战：

- 数据质量：质量差的数据可能导致不准确的模型。
- 过度拟合：当模型过度学习训练数据，导致在新数据上表现不佳。
- 偏见和公平性：模型可能会延续或放大数据中的现有偏见。
- 可解释性：复杂的模型如深度神经网络可能难以解释。
- 计算资源：训练大型模型可能需要大量的计算能力。

机器学习的未来
--------------------------

机器学习领域正在迅速发展，趋势包括：

- 自动机器学习(AutoML)：自动化应用机器学习过程的工具。
- 联邦学习：在去中心化设备上训练模型，同时保持数据本地化。
- 边缘AI：在智能手机和物联网设备等边缘设备上部署机器学习模型。
- 可解释AI：开发使机器学习模型更具可解释性的技术。

结论
----------

机器学习是一个强大的工具，正在改变许多行业。随着该领域的不断发展，
从业者了解这些技术的技术方面和伦理影响都很重要。
"""

    # 创建示例文档
    os.makedirs("sample_docs", exist_ok=True)
    doc_path = "sample_docs/machine_learning_intro.txt"
    
    with open(doc_path, "w", encoding="utf-8") as f:
        f.write(sample_doc.strip())
    
    return doc_path


async def main_async():
    parser = argparse.ArgumentParser(
        description="使用合成数据生成技术从文档中生成高质量的问答对"
    )
    parser.add_argument(
        "--document",
        default="/Users/sd/Desktop/mycode/myalgo/milvus/ai-insurance-expert/badcase/zcz/like.txt",
        help="输入文档的路径。如果未提供，将创建并使用示例文档。"
    )
    parser.add_argument(
        "--num-pairs",
        type=int, 
        default=20,
        help="每个文本块生成的问答对数量"
    )
    parser.add_argument(
        "--chunk-size", 
        type=int, 
        default=800, 
        help="每个文本块的大小"
    )
    parser.add_argument(
        "--overlap", 
        type=int, 
        default=100, 
        help="文本块之间的重叠"
    )
    parser.add_argument(
        "--output-folder", 
        default="generated_data", 
        help="生成数据的输出文件夹"
    )
    parser.add_argument(
        "--model",
        default="qwq-32b-preview",
        help="用于问答生成的模型"
    )
    parser.add_argument(
        "--async-mode",
        default=True,
        action=argparse.BooleanOptionalAction,
        help="使用异步模式生成数据（默认：开启）"
    )

    args = parser.parse_args()

    # 如果没有提供文档，则创建一个示例文档
    if args.document:
        document_path = args.document
    else:
        document_path = create_sample_document()
        print(f"未提供文档。在以下位置创建了一个示例文档: {document_path}")

    # 创建生成器
    generator = SyntheticDataGenerator(output_folder=args.output_folder, model=args.model)

    print("开始生成合成问答对...")
    print(f"文档: {document_path}")
    print(f"每个文本块的问答对数量: {args.num_pairs}")
    print(f"文本块大小: {args.chunk_size}")
    print(f"重叠: {args.overlap}")
    print(f"输出文件夹: {args.output_folder}")
    print(f"模型: {args.model}")
    print(f"异步模式: {args.async_mode}")
    print("-" * 50)

    # 根据模式选择同步或异步生成数据集
    if args.async_mode:
        final_files = await generator.generate_dataset_async(
            document_path=document_path,
            num_pairs=args.num_pairs,
            chunk_size=args.chunk_size,
            overlap=args.overlap
        )
    else:
        final_files = generator.generate_dataset(
            document_path=document_path,
            num_pairs=args.num_pairs,
            chunk_size=args.chunk_size,
            overlap=args.overlap
        )

    print("-" * 50)
    print("数据集生成成功完成！")
    print("最终生成的文件:")
    for file in final_files:
        print(f"  - {file}")

    print("\n要检查生成的数据，可以使用:")
    print("  cat generated_data/final/*.json")


def main():
    parser = argparse.ArgumentParser(
        description="使用合成数据生成技术从文档中生成高质量的问答对"
    )
    parser.add_argument(
        "--document",
        default="/Users/sd/Desktop/mycode/myalgo/milvus/ai-insurance-expert/badcase/zcz/like.txt",
        help="输入文档的路径。如果未提供，将创建并使用示例文档。"
    )
    parser.add_argument(
        "--num-pairs",
        type=int, 
        default=20,
        help="每个文本块生成的问答对数量"
    )
    parser.add_argument(
        "--chunk-size", 
        type=int, 
        default=800, 
        help="每个文本块的大小"
    )
    parser.add_argument(
        "--overlap", 
        type=int, 
        default=100, 
        help="文本块之间的重叠"
    )
    parser.add_argument(
        "--output-folder", 
        default="generated_data", 
        help="生成数据的输出文件夹"
    )
    parser.add_argument(
        "--model",
        default="qwq-32b-preview",
        help="用于问答生成的模型"
    )
    parser.add_argument(
        "--async-mode",
        default=True,
        action=argparse.BooleanOptionalAction,
        help="使用异步模式生成数据（默认：开启）"
    )

    args = parser.parse_args()
    
    # 强制设置异步模式为True（如果需要）
    args.async_mode = True

    # 如果没有提供文档，则创建一个示例文档
    if args.document:
        document_path = args.document
    else:
        document_path = create_sample_document()
        print(f"未提供文档。在以下位置创建了一个示例文档: {document_path}")

    # 创建生成器
    generator = SyntheticDataGenerator(output_folder=args.output_folder, model=args.model)

    print("开始生成合成问答对...")
    print(f"文档: {document_path}")
    print(f"每个文本块的问答对数量: {args.num_pairs}")
    print(f"文本块大小: {args.chunk_size}")
    print(f"重叠: {args.overlap}")
    print(f"输出文件夹: {args.output_folder}")
    print(f"模型: {args.model}")
    print(f"异步模式: {args.async_mode}")
    print("-" * 50)

    # 根据模式选择同步或异步生成数据集
    if args.async_mode:
        final_files = asyncio.run(generator.generate_dataset_async(
            document_path=document_path,
            num_pairs=args.num_pairs,
            chunk_size=args.chunk_size,
            overlap=args.overlap
        ))
    else:
        final_files = generator.generate_dataset(
            document_path=document_path,
            num_pairs=args.num_pairs,
            chunk_size=args.chunk_size,
            overlap=args.overlap
        )

    print("-" * 50)
    print("数据集生成成功完成！")
    print("最终生成的文件:")
    for file in final_files:
        print(f"  - {file}")

    print("\n要检查生成的数据，可以使用:")
    print("  cat generated_data/final/*.json")


if __name__ == '__main__':
    main()